import os
import time
import datetime
from langgraph.graph import StateGraph, END
# from langgraph.checkpoint.memory import MemorySaver
from .utils.views import print_agent_output
from ..memory.research import ResearchState
from .utils.utils import sanitize_filename

# Import agent classes
from . import \
    WriterAgent, \
    EditorAgent, \
    PublisherAgent, \
    ResearchAgent, \
    HumanAgent



"""
    面向对象编程
        1. 编排类统一管理多Agent。
        2. 多Agent 和上层有明显的分层。
"""
class ChiefEditorAgent:
    """
    主编辑器代理 - 负责管理和协调研究任务的执行流程。
    
    该类作为一个编排者(Orchestrator)，负责初始化其他代理(如研究员、编辑、写手等)，
    创建它们之间的工作流，并协调整个研究过程的执行。它管理任务ID、输出目录、
    以及代理之间的消息传递和工作流转换。
    
    Agent responsible for managing and coordinating editing tasks.
    """

    def __init__(self, task: dict, websocket=None, stream_output=None, tone=None, headers=None):
        """
        初始化主编辑器代理。
        
        参数:
            task (dict): 研究任务的详细信息
            websocket: WebSocket连接对象，用于实时通信
            stream_output: 用于流式输出的函数
            tone: 输出的语气风格
            headers: HTTP请求头信息
        """
        self.task = task
        self.websocket = websocket
        self.stream_output = stream_output
        self.headers = headers or {}
        self.tone = tone
        self.task_id = self._generate_task_id()
        self.output_dir = self._create_output_directory()

    def _generate_task_id(self):
        """
        生成唯一的任务ID。
        
        目前基于时间戳生成，但可以是任何唯一标识符。
        
        返回:
            int: 唯一的任务ID
        """
        # 目前基于时间戳生成，但可以是任何唯一标识符
        return int(time.time())

    def _create_output_directory(self):
        """
        创建输出目录用于存储研究结果。
        
        目录名基于任务ID和查询内容生成。
        
        返回:
            str: 输出目录的路径
        """
        output_dir = "./outputs/" + \
            sanitize_filename(
                f"run_{self.task_id}_{self.task.get('query')[0:40]}")

        os.makedirs(output_dir, exist_ok=True)
        return output_dir

    def _initialize_agents(self):
        """
        初始化所有参与研究过程的代理。
        
        包括写手、编辑、研究员、发布者和人类代理。
        
        返回:
            dict: 包含所有初始化代理的字典
        """
        return {
            "writer": WriterAgent(self.websocket, self.stream_output, self.headers),
            "editor": EditorAgent(self.websocket, self.stream_output, self.tone, self.headers),
            "research": ResearchAgent(self.websocket, self.stream_output, self.tone, self.headers),
            "publisher": PublisherAgent(self.output_dir, self.websocket, self.stream_output, self.headers),
            "human": HumanAgent(self.websocket, self.stream_output, self.headers)
        }

    def _create_workflow(self, agents):
        """
        创建代理之间的工作流程图。
        
        为每个代理添加节点，并定义它们之间的工作流转换。
        
        参数:
            agents (dict): 包含所有代理的字典
            
        返回:
            StateGraph: 配置好的工作流程图
        """
        workflow = StateGraph(ResearchState)

        # Add nodes for each agent
        workflow.add_node("browser", agents["research"].run_initial_research)
        workflow.add_node("planner", agents["editor"].plan_research)
        workflow.add_node("researcher", agents["editor"].run_parallel_research)
        workflow.add_node("writer", agents["writer"].run)
        workflow.add_node("publisher", agents["publisher"].run)
        workflow.add_node("human", agents["human"].review_plan)

        # Add edges
        self._add_workflow_edges(workflow)

        return workflow

    def _add_workflow_edges(self, workflow):
        """
        添加工作流程图中节点之间的连接关系。
        
        定义代理之间的工作流程转换，包括条件转换逻辑。
        
        参数:
            workflow (StateGraph): 工作流程图对象
        """
        workflow.add_edge('browser', 'planner')
        workflow.add_edge('planner', 'human')
        workflow.add_edge('researcher', 'writer')
        workflow.add_edge('writer', 'publisher')
        workflow.set_entry_point("browser")
        workflow.add_edge('publisher', END)

        # Add human in the loop
        workflow.add_conditional_edges(
            'human',
            lambda review: "accept" if review['human_feedback'] is None else "revise",
            {"accept": "researcher", "revise": "planner"}
        )

    def init_research_team(self):
        """
        初始化研究团队并创建其工作流程。
        
        返回:
            StateGraph: 配置好的研究团队工作流程图
        """
        agents = self._initialize_agents()
        return self._create_workflow(agents)

    async def _log_research_start(self):
        """
        记录研究任务的开始。
        
        根据可用的输出方式（WebSocket或控制台）输出开始信息。
        """
        message = f"Starting the research process for query '{self.task.get('query')}'..."
        if self.websocket and self.stream_output:
            await self.stream_output("logs", "starting_research", message, self.websocket)
        else:
            print_agent_output(message, "MASTER")

    async def run_research_task(self, task_id=None):
        """
        运行研究任务，协调整个研究团队完成工作。
        
        参数:
            task_id (optional): 任务的ID
            
        返回:
            研究任务的结果
        """
        research_team = self.init_research_team()
        chain = research_team.compile()

        await self._log_research_start()

        config = {
            "configurable": {
                "thread_id": task_id,
                "thread_ts": datetime.datetime.utcnow()
            }
        }

        result = await chain.ainvoke({"task": self.task}, config=config)
        return result
